# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import functools
from collections import namedtuple, OrderedDict

from coin.base.datetime_util import to_datetime, to_timestamp_int
from coin.exchange.base.order_gateway import OrderSide
from coin.exchange.kr_rest.product_holders import get_holder_from_product

from coin.strategy.hk_hard_arb.index_pricer import IndexPricer
from coin.strategy.mm.simple_sim.executors import PassiveSimExecutor
from coin.strategy.mm.subscription import get_subscription_key

FEATURE_ROW_FIELDS = [
    'time',
    'time_hr',
    'ask0_1',
    'askt_1',
    'bid0_1',
    'bidt_1',
    'ask0_2',
    'askt_2',
    'bid0_2',
    'bidt_2',
    'ask_basis_ma',
    'bid_basis_ma'
]
_FeatureRow = namedtuple('FeatureRow', FEATURE_ROW_FIELDS)


class FeatureRow(_FeatureRow):
  @staticmethod
  def get_csv_header():
    return 'header_feature_dump,%s' % ','.join(_FeatureRow._fields)

  def to_csv(self):
    return ('"FEATURE_DUMP",%d,"%s",'
            '%f,%f,%f,%f,%f,%f,%f,%f,%f,%f' % (self.time,
                                               self.time_hr,
                                               self.ask0_1,
                                               self.askt_1,
                                               self.bid0_1,
                                               self.bidt_1,
                                               self.ask0_2,
                                               self.askt_2,
                                               self.bid0_2,
                                               self.bidt_2,
                                               self.ask_basis_ma,
                                               self.bid_basis_ma))

  def to_log_string(self):
    return '\n' + self.to_csv()


FILL_ROW_FIELDS = ['time', 'time_hr', 'side', 'price', 'qty', 'maker', 'pos', 'balance']
_FillRow = namedtuple('FillRow', FILL_ROW_FIELDS)


class FillRow(_FillRow):
  @staticmethod
  def get_csv_header():
    return 'header_fill_dump,%s' % ','.join(FILL_ROW_FIELDS)

  def to_csv(self):
    return ('"FILL_DUMP",%d,"%s",'
            '%d,%f,%f,%d,%f,%f' % (self.time,
                                   self.time_hr,
                                   self.side,
                                   self.price,
                                   self.qty,
                                   self.maker,
                                   self.pos,
                                   self.balance))

  def to_log_string(self):
    return '\n' + self.to_csv()


class PassUnhedgedSimStrategy(object):
  def __init__(
      self,
      trade_product,
      ref_products,
      sell_edge_func,  # pos --> sell_edge
      buy_edge_func,  # pos --> buy_edge
      pricing_config,
      executor_param,
      agg_edge=None,  # If None, does not aggress.
      trade_after=0,
      name=None,
      feature_filepath=None,
      fill_filepath=None,
      bbo_protection=False,
      pull_edge=None):
    self._trade_product = trade_product
    assert isinstance(ref_products, list)
    self._ref_products = ref_products

    self._trade_after = to_timestamp_int(trade_after)
    self._start_ts = None

    self._sell_edge_func = sell_edge_func
    self._buy_edge_func = buy_edge_func
    self._agg_edge = agg_edge
    tick = get_holder_from_product(trade_product).product_info.gen_tick()
    self._pricer = IndexPricer(pricing_config, tick)
    self._bbo_protection = bbo_protection
    self._pull_edge = pull_edge

    do_agg = agg_edge is not None

    unavail_intervals_file = None
    if trade_product.exchange == 'Bitmex':
      unavail_intervals_file = 'coin/strategy/data/bitmex_unavailability.txt'

    self._executor = PassiveSimExecutor(do_agg=do_agg,
                                        unavail_intervals_file=unavail_intervals_file,
                                        **executor_param)
    self._executor.pos_callback = self.on_pos_update
    self._trade_book = None

    self.name = name
    self._feature_filepath = feature_filepath
    self._feature_file = None
    self._last_feature_dump_ts = 0
    self._fill_filepath = fill_filepath
    self._fill_file = None
    self.fills = []
    self._ref_books = OrderedDict()
    for ref_product in ref_products:
      self._ref_books[ref_product] = None

  @property
  def trade_product(self):
    return self._trade_product

  @property
  def book_ready(self):
    return self._trade_book is not None and all(
        [ref_book is not None for ref_book in self._ref_books.values()])

  @property
  def ref_books_list(self):
    return list(self._ref_books.values())

  def _add_book_builder_callback(self, book_builder, product):
    book_builder.subscribe(product, functools.partial(self.on_book, product))
    if hasattr(book_builder, 'subscribe_trade'):
      book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))
    if hasattr(book_builder, "set_parse_only_subscribed"):
      book_builder.set_parse_only_subscribed()

  def on_book_reset(self, book_builder_name, book_builder):
    for product in set([self._trade_product] + self._ref_products):
      subscription_key = get_subscription_key(product)
      if not book_builder_name == subscription_key:
        continue
      self._add_book_builder_callback(book_builder, product)

  def on_book(self, product, book):
    is_trade_product = (product == self._trade_product)
    if is_trade_product:
      self._executor.update_book(book)
      self._trade_book = book

    if product in self._ref_books:
      self._ref_books[product] = book

    if not self.book_ready:
      return

    if self._start_ts is None:
      self._start_ts = book.timestamp + self._trade_after

    self._pricer.update(self.ref_books_list, self._trade_book)

    if not self._pricer.ready or book.timestamp < self._start_ts:
      return

    sell_edge = self._sell_edge_func(self._executor.pos.pos)
    buy_edge = self._buy_edge_func(self._executor.pos.pos)
    sell_price_pull = None
    buy_price_pull = None
    try:
      pass_sell_p, pass_buy_p = self._pricer.calculate_price(sell_edge,
                                                             buy_edge,
                                                             bbo_protection=self._bbo_protection)
      if pass_sell_p is not None and self._pull_edge is not None:
        sell_price_pull = pass_sell_p * (1 - self._pull_edge)
      if pass_buy_p is not None and self._pull_edge is not None:
        buy_price_pull = pass_buy_p * (1 + self._pull_edge)
    except Exception as e:
      print('Fail to calculate price.', type(e), e)
      return

    if self._agg_edge is not None:
      self._executor.cancel_agg_orders(book.timestamp)
      agg_sell_p, agg_buy_p = self._pricer.calculate_price(self._agg_edge, self._agg_edge, False)
      if agg_sell_p <= pass_sell_p:
        self._executor.agg_sell(agg_sell_p, book.timestamp)
      if agg_buy_p >= pass_buy_p:
        self._executor.agg_buy(agg_buy_p, book.timestamp)

    self._executor.post(timestamp=book.timestamp,
                        sell_price=pass_sell_p,
                        buy_price=pass_buy_p,
                        sell_price_pull=sell_price_pull,
                        buy_price_pull=buy_price_pull)

    self._dump_feature()

  def on_trade(self, product, trade):
    if product == self._trade_product:
      assert trade.symbol == self._trade_product.native_symbol
      self._executor.update_trade(trade)

  @property
  def pos(self):
    return self._executor.pos

  def get_summary(self):
    return self._executor.get_summary()

  def __del__(self):
    if self._feature_file is not None:
      self._feature_file.close()
    if self._fill_file is not None:
      self._fill_file.close()

  @property
  def _feature_dump_enabled(self):
    return self._feature_filepath is not None

  def _ensure_feature_file(self):
    if self._feature_dump_enabled and self._feature_file is None:
      self._feature_file = open(self._feature_filepath, 'w')
      self._feature_file.write(FeatureRow.get_csv_header() + '\n')

  def _dump_feature(self):
    if not self._feature_dump_enabled:
      return

    row = self._pricer.last_value
    if row is None or row.time - self._last_feature_dump_ts < 2 * (10**9):
      return
    self._last_feature_dump_ts = row.time

    self._ensure_feature_file()
    self._feature_file.write(row.to_csv() + '\n')

  @property
  def _fill_dump_enabled(self):
    return self._fill_filepath is not None

  def _ensure_fill_file(self):
    if self._fill_dump_enabled and self._fill_file is None:
      self._fill_file = open(self._fill_filepath, 'w')
      self._fill_file.write(FillRow.get_csv_header() + '\n')

  def on_pos_update(self, pos, price, qty, side, maker, timestamp):
    side = 1 if side == OrderSide.BUY else -1
    row = FillRow(timestamp, str(to_datetime(timestamp)), side, price, qty, maker, pos.pos, pos.pnl)
    self.fills.append(row)

    if not self._fill_dump_enabled:
      return

    self._ensure_fill_file()
    self._fill_file.write(row.to_csv() + '\n')
